package tagdb.algorithm.tag;

import tagdb.denovo.DenovoHit;
import tagdb.io.MascotGenericFile;

import java.util.ArrayList;
import java.util.List;
import java.util.Set;
import java.util.TreeSet;

/**
 * <b>TagGenerator</b>
 * <p>
 * This class provides functions to generate a set of N-terminal and C-terminal sequence tags.
 * As input a set of previously calculated denovo sequences is needed.
 * </p>
 *
 * @author Thilo Muth
 */
public class TagGenerator {

    /**
     * The list of denovo hits.
     */
    private List<DenovoHit> hits;

    /**
     * Holds the minimum N-terminal tag length.
     */
    private int nTagLength;

    /**
     * Holds the minimum C-terminal tag length.
     */
    private int cTagLength;

    /**
     * Spectrum object.
     */
    private MascotGenericFile spectrum;

    /**
     * Constructs the TagGenerator object, it takes a list of denovo hits as seed.
     *
     * @param hits       List of denovo hits as seed.
     * @param spectrum   The spectrum.
     * @param nTagLength Length of the N-terminal tag.
     * @param cTagLength Length of the C-terminal tag.
     */
    public TagGenerator(List<DenovoHit> hits, MascotGenericFile spectrum, int nTagLength, int cTagLength) {
        this.hits = hits;
        this.spectrum = spectrum;
        this.nTagLength = nTagLength;
        this.cTagLength = cTagLength;
    }

    /**
     * Returns all tags.
     * @return all tags
     */
    public List<Tag> getAllTags() {
        List<Tag> tags = new ArrayList<Tag>();
        tags.addAll(getCTermTags());
        tags.addAll(getNTermTags());
        return tags;
    }

    /**
     * Generates a list of C-terminal tags.
     *
     * @return List of C-terminal tags.
     */
    public List<Tag> getCTermTags() {
        // Set of tags.
        Set<Tag> tags = new TreeSet<Tag>();

        // Get C-Terminal tags
        for (DenovoHit hit : hits) {
            if (hit.getCTermGap() == 0.0) {

                // Minimum tag length.
                if(hit.getLength() < cTagLength) continue;
                String seq = hit.getSequence().substring(hit.getLength() - cTagLength, hit.getLength());

                // Create tag
                Tag tag = new Tag(seq, Tag.TagType.C_TERMINAL, hit.getCharge());
                new TagScorer(tag, spectrum);

                tags.add(tag);

                Set<Tag> errorTolTags = getErrorTolerantCTags(tag);
                for(Tag errorTag : errorTolTags) {
                    new TagScorer(errorTag, spectrum);
                    tags.add(errorTag);
                }
            }
        }
        return new ArrayList<Tag>(tags);
    }

    /**
     * Generates a list of C-terminal tags.
     *
     * @return List of C-terminal tags.
     */
    public List<Tag> getNTermTags() {
        // Set of tags.
        Set<Tag> tags = new TreeSet<Tag>();

        // Get C-Terminal tags
        for (DenovoHit hit : hits) {
            if (hit.getNTermGap() == 0.0) {

                // Minimum tag length.
                if(hit.getLength() < nTagLength) continue;
                String seq = hit.getSequence().substring(0, nTagLength);

                // Create tag
                Tag tag = new Tag(seq, Tag.TagType.N_TERMINAL, hit.getCharge());
                new TagScorer(tag, spectrum);
                tags.add(tag);

                Set<Tag> errorTolTags = getErrorTolerantNTags(tag);
                for(Tag errorTag : errorTolTags) {
                    new TagScorer(errorTag, spectrum);
                    tags.add(errorTag);
                }
            }
        }
        return new ArrayList<Tag>(tags);
    }

    /**
     * Calculates the error-tolerant tags for a given tag.
     * @param tag Tag given.
     * @return Set of error-tolerant tags.
     */
    public Set<Tag> getErrorTolerantCTags(Tag tag) {
        // Account for sequencing errors
        List<String> errorSeqs = new ArrayList<String>();
        Set<Tag> errorTolTags = new TreeSet<Tag>();

        String tagSeq = tag.getSequence().substring(0, cTagLength - 1);
        String endSeq = tag.getSequence().substring(cTagLength - 1, cTagLength);
        if (tagSeq.contains("Q")) {
            errorSeqs.add(tagSeq.replaceAll("Q", "K") + endSeq);
        }
        if (tagSeq.contains("K")) {
                errorSeqs.add(tagSeq.replaceAll("K", "Q") + endSeq);
        }
        if (tagSeq.contains("I")) {
            errorSeqs.add(tagSeq.replaceAll("I", "L") + endSeq);
        }
        if (tagSeq.contains("II")) {
            errorSeqs.add(tagSeq.replaceAll("II", "LL") + endSeq);
        }
        if (tagSeq.contains("LL")) {
            errorSeqs.add(tagSeq.replaceAll("LL", "II") + endSeq);
        }
        if (tagSeq.contains("L")) {
            errorSeqs.add(tagSeq.replaceAll("L", "I") + endSeq);
        }
        if (tagSeq.contains("L") && tagSeq.contains("Q")) {
            String errorSeq = tagSeq.replaceFirst("L", "I");
            errorSeqs.add(errorSeq.replaceFirst("Q", "K") + endSeq);
        }

        // Add the tag itself to the error sequences for following replacements.
        errorSeqs.add(tag.getSequence());

        // Multiple amino acid replacements.
        for (String errorSeq : errorSeqs) {
            if (errorSeq.indexOf("KR") != -1) {
                errorTolTags.add(new Tag(errorSeq.replaceAll("KR", "RK"), tag.getType(), tag.getCharge()));
            }
            if (errorSeq.contains("RK")) {
                errorTolTags.add(new Tag(errorSeq.replaceAll("RK", "KR"), tag.getType(), tag.getCharge()));
            }
            if (errorSeq.contains("RK")) {
                errorTolTags.add(new Tag(errorSeq.replaceAll("RK", "QR"), tag.getType(), tag.getCharge()));
            }
            if (errorSeq.contains("AR")) {
                errorTolTags.add(new Tag(errorSeq.replaceAll("AR", "KV"), tag.getType(), tag.getCharge()));
            }
            if (errorSeq.contains("KV")) {
                errorTolTags.add(new Tag(errorSeq.replaceAll("KV", "AR"), tag.getType(), tag.getCharge()));
            }
            if (errorSeq.contains("MP")) {
                errorTolTags.add(new Tag(errorSeq.replaceAll("MP", "EV"), tag.getType(), tag.getCharge()));
            }
            if (errorSeq.contains("EV")) {
                errorTolTags.add(new Tag(errorSeq.replaceAll("EV", "MP"), tag.getType(), tag.getCharge()));
            }
            if (errorSeq.contains("HFA")) {
                errorTolTags.add(new Tag(errorSeq.replaceAll("HFA", "PEE"), tag.getType(), tag.getCharge()));
            }
            if (errorSeq.contains("PEE")) {
                errorTolTags.add(new Tag(errorSeq.replaceAll("PEE", "HFA"), tag.getType(), tag.getCharge()));
            }
            errorTolTags.add(new Tag(errorSeq, tag.getType(), tag.getCharge()));
        }
        return errorTolTags;
    }

    /**
     * Calculates the error-tolerant tags for a given tag.
     * @param tag Tag given.
     * @return Set of error-tolerant tags.
     */
    public Set<Tag> getErrorTolerantNTags(Tag tag) {
        // Account for sequencing errors
        List<String> errorSeqs = new ArrayList<String>();
        Set<Tag> errorTolTags = new TreeSet<Tag>();

        String tagSeq = tag.getSequence();

        if (tagSeq.contains("Q")) {
            errorSeqs.add(tagSeq.replaceAll("Q", "K"));
        }
        if (tagSeq.contains("K")) {
            errorSeqs.add(tagSeq.replaceAll("K", "Q"));
        }
        if (tagSeq.contains("I")) {
            errorSeqs.add(tagSeq.replaceAll("I", "L"));
        }
        if (tagSeq.contains("L")) {
            errorSeqs.add(tagSeq.replaceAll("L", "I"));
        }
        if (tagSeq.contains("II")) {
            errorSeqs.add(tagSeq.replaceAll("II", "LL") );
        }
        if (tagSeq.contains("LL")) {
            errorSeqs.add(tagSeq.replaceAll("LL", "II"));
        }

        // Add the tag itself to the error sequences for following replacements.
        errorSeqs.add(tag.getSequence());

        // Multiple amino acid replacements.
        for (String errorSeq : errorSeqs) {
            if (errorSeq.contains("MP")) {
                errorTolTags.add(new Tag(errorSeq.replaceAll("MP", "EV"), tag.getType(), tag.getCharge()));
            }
            if (errorSeq.contains("NLE")) {
                errorTolTags.add(new Tag(errorSeq.replaceAll("NLE", "DQI"), tag.getType(), tag.getCharge()));
            }
            if (errorSeq.contains("NLE")) {
                errorTolTags.add(new Tag(errorSeq.replaceAll("NLE", "DQL"), tag.getType(), tag.getCharge()));
            }
            if (errorSeq.contains("HFA")) {
                errorTolTags.add(new Tag(errorSeq.replaceAll("HFA", "PEE"), tag.getType(), tag.getCharge()));
            }
            if (errorSeq.contains("PEE")) {
                errorTolTags.add(new Tag(errorSeq.replaceAll("PEE", "HFA"), tag.getType(), tag.getCharge()));
            }
            errorTolTags.add(new Tag(errorSeq, tag.getType(), tag.getCharge()));
        }
        return errorTolTags;
    }
}